#include "bs.h"
#include <cstdio>
#include <cstdlib>

void setup_cuda(BS_test_t *t);
void setup_buffers(BS_test_t * t, float(*gen)(float low, float high));
void warmup(BS_test_t * t, int num_iters);
bool run(BS_test_t * t, int num_iters);
bool check(BS_test_t *t);
	void BS_gold(BS_test_t *t);
		extern "C" void BlackScholesCPU(
		float *h_CallResult,
		float *h_PutResult,
		float *h_StockPrice,
		float *h_OptionStrike,
		float *h_OptionYears,
		float Riskfree,
		float Volatility,
		int optN
	);
void print_arr(BS_test_t * t, float *arr);
void post(BS_test_t *t, int run_num_iters, bool result);
void cleanup(BS_test_t *t);

/////////////////////////////////////////////////////

    const unsigned int   optionCount = 4000000;
    const float                    R = 0.02f;
    const float                    V = 0.30f;

/////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////
// Helper function, returning uniformly distributed
// random float in [low, high] range
// Copyright 1993-2010 NVIDIA Corporation.  All rights reserved.
////////////////////////////////////////////////////////////////////////////////
float RandFloat(float low, float high){
    float t = (float)rand() / (float)RAND_MAX;
    return (1.0f - t) * low + t * high;
}

void setup_buffers(BS_test_t * t, float(*gen)(float low, float high))
{
	// malloc host memory
	t->h_CallResult = (float *) malloc (t->num * sizeof(float));
	t->h_PutResult = (float *) malloc (t->num * sizeof(float));
	t->h_StockPrice = (float *) malloc (t->num * sizeof(float));
	t->h_OptionStrike = (float *) malloc (t->num * sizeof(float));
	t->h_OptionYears = (float *) malloc (t->num * sizeof(float));

	// create GPU buffers
	printf("Creating GPU buffers...\n");
	cutilSafeCall(cudaMalloc(&(t->d_CallResult), t->num * sizeof(float)));
	cutilSafeCall(cudaMalloc(&(t->d_PutResult), t->num * sizeof(float)));
	cutilSafeCall(cudaMalloc(&(t->d_StockPrice), t->num * sizeof(float)));
	cutilSafeCall(cudaMalloc(&(t->d_OptionStrike), t->num * sizeof(float)));
	cutilSafeCall(cudaMalloc(&(t->d_OptionYears), t->num * sizeof(float)));

	// init host memory
	for (unsigned int i = 0; i < t->num; i++) {
		t->h_CallResult[i] = -1.0f;
		t->h_PutResult[i]  = -1.0f;
		t->h_StockPrice[i] = gen(5.0f, 30.0f);
		t->h_OptionStrike[i] = gen(1.0f, 100.0f);
		t->h_OptionYears[i] = gen(0.25f, 10.0f);
	}

		// copy to device memory
	printf("Copying input matricies to GPU...\n");
	cutilSafeCall(cudaMemcpy(t->d_StockPrice, t->h_StockPrice, t->num * sizeof(float), cudaMemcpyHostToDevice));
	cutilSafeCall(cudaMemcpy(t->d_OptionStrike, t->h_OptionStrike, t->num * sizeof(float), cudaMemcpyHostToDevice));
	cutilSafeCall(cudaMemcpy(t->d_OptionYears, t->h_OptionYears, t->num * sizeof(float), cudaMemcpyHostToDevice));
}

void setup_cuda(BS_test_t *t)
{
	// TODO
} 

float gen0(float low, float high) {
	srand(5347);
	return RandFloat(low, high);
}

float gen1(float low, float high) {
	return 5.0f;
}

void warmup(BS_test_t * t, int num_iters)
{
	for (int i = 0; i < num_iters; i++) {
		launch_BS(t);
		printf("Warmup %d/%d complete\n", i+1, num_iters);
	}
}

bool run(BS_test_t * t, int num_iters)
{
	t->time = 0.0;

	for (int i = 0; i < num_iters; i++) {
		double time = launch_BS(t);
		t->time += time;
		printf("Run %d complete\n", i+1, num_iters);
	}
	if (!check(t)) 
		return false;
	return true;
}

bool check(BS_test_t *t)
{
	float EPSILON = 2.0f;

	BS_gold(t);
	float * CallResult = (float *) malloc (t->num * sizeof(float));
	float * PutResult = (float *) malloc (t->num * sizeof(float));
	cutilSafeCall(cudaMemcpy(CallResult, t->d_CallResult, t->num * sizeof(float),
		cudaMemcpyDeviceToHost));
	cutilSafeCall(cudaMemcpy(PutResult, t->d_PutResult, t->num * sizeof(float),
		cudaMemcpyDeviceToHost));
	for (unsigned int i = 0; i < t->num; i++) {
		float hstCall = t->h_CallResult[i];
		float hstPut = t->h_PutResult[i];
		float devCall = CallResult[i];
		float devPut = PutResult[i];
		if (abs(hstCall - devCall) >= EPSILON ||
			abs(hstPut - devPut) >= EPSILON ) {
				fprintf(stderr, "check failed at %d: [%f, %f] != [%f, %f]\n", 
					i,
					hstCall, hstPut,
					devCall, devPut);
				free(CallResult);
				free(PutResult);
				return false;
		}
	}
	free(CallResult);
	free(PutResult);
	return true;
}

void BS_gold(BS_test_t * t)
{
	BlackScholesCPU(
		t->h_CallResult,
		t->h_PutResult,
		t->h_StockPrice,
		t->h_OptionStrike,
		t->h_OptionYears,
		t->Riskfree,
		t->Volatility,
		t->num
	);
}

void print_arr(BS_test_t * t, float *arr)
{
	for (unsigned int i = 0; i < t->num; i++) {
		if (i == 0)
			printf("[");
		printf("%3.0f", arr[i]);
		if (i < t->num-1)
			printf(", ");
		else
			printf("]\n");
	}
}

void post(BS_test_t * t, int run_num_iters, bool result)
{
	printf("----------------RESULTS----------------\n");
	printf("PASS? %s\n", result ? "YES" : "NO");
	printf("Size: %d threads, block size of %d\n", t->gridSize, t->blockSize);
	printf("Time: %f\n", t->time / run_num_iters);
}

void output(BS_test_t *t, FILE *csv, int run_num_iters) {
	fprintf(csv, "%u, %u, %u, %1.9Lf\n",
		t->num,
		t->blockSize,
		t->gridSize,
		t->time / run_num_iters);
}

void cleanup(BS_test_t * t)
{
	printf("Freeing Host and Device buffers...\n");
	free(t->h_CallResult);
	free(t->h_PutResult);
	free(t->h_StockPrice);
	free(t->h_OptionStrike);
	free(t->h_OptionYears);

	cutilSafeCall(cudaFree(t->d_CallResult));
	cutilSafeCall(cudaFree(t->d_PutResult));
	cutilSafeCall(cudaFree(t->d_StockPrice));
	cutilSafeCall(cudaFree(t->d_OptionStrike));
	cutilSafeCall(cudaFree(t->d_OptionYears));
}

void test_BS (unsigned int sizeMin, unsigned int sizeMax, 
			  int warmup_num_iters, int run_num_iters, 
			  int blockMin, int blockMax, 
			  int gridMin, int gridMax,
			  char *filename, float(*gen)(float low, float high))
{
	BS_test_t t;
	t.num = sizeMin;
	t.blockSize = blockMin;
	t.gridSize = gridMin;
	t.Riskfree = R;
	t.Volatility = V;
	bool result = false;

	FILE * csv = fopen(filename, "a");
	if (csv == NULL)
		fprintf(stderr, "can't open output file\n");

	if (csv != NULL)
		fprintf(csv, "size, block size, grid size, total time\n");

	setup_cuda(&t);
	setup_buffers(&t, gen);

	warmup(&t, warmup_num_iters);

	for (t.num = sizeMin; t.num <= sizeMax; t.num <<= 1) {
		for (t.gridSize = gridMin; t.gridSize <= gridMax; t.gridSize += 16*1024) {
			for (t.blockSize = blockMin; t.blockSize <= blockMax; t.blockSize += 64) {
				if (t.gridSize % t.blockSize != 0)
					continue;
				setup_buffers(&t, gen);
				result = run(&t, run_num_iters);
				post(&t, run_num_iters, result);
				if (result && csv != NULL)
					output(&t, csv, run_num_iters);
				cleanup(&t);
			}
		}
	}

	if (csv != NULL)
		fclose(csv);
}

int main(int argc, char *argv[])
{
	char *filename = "bs_cuda.csv";
	test_BS(optionCount, optionCount, 5, 100, 64, 512, 16*1024, 64*1024, filename, gen0);
	return 0;
}
